set_project("AIcompiler")
set_version("0.1.0")
set_languages("c++17")

-- 选项
option("with_cuda")
    set_default(false)
    set_showmenu(true)
    set_description("Enable CUDA/PTX backend (later phase)")
option_end()

option("with_pocl")
    set_default(true)
    set_showmenu(true)
    set_description("Enable PoCL/SPIR-V backend")
option_end()

-- 全局包含与定义
add_includedirs("runtime/common", "mlir/dialects")
add_defines("AICOMPILER_VERSION=\"0.1.0\"")

-- LLVM/MLIR 路径检测
local llvm_dir = os.getenv("LLVM_DIR")
if llvm_dir then
    add_includedirs(path.join(llvm_dir, "../../../include"))
    add_linkdirs(path.join(llvm_dir, "../../../lib"))
elseif os.isdir("third_party/llvm-install") then
    add_includedirs("third_party/llvm-install/include")
    add_linkdirs("third_party/llvm-install/lib")
end

-- 运行时库：CPU
target("aic_runtime_cpu")
    set_kind("static")
    add_files("runtime/cpu/*.cpp")
    add_headerfiles("runtime/common/*.h")

-- 运行时库：SPIR-V/PoCL（占位）
target("aic_runtime_spirv")
    set_kind("static")
    add_files("runtime/spirv/*.cpp")
    add_headerfiles("runtime/common/*.h")
    add_defines("AIC_WITH_POCL")

-- TableGen 自定义规则
rule("mlir_tablegen")
    set_extensions(".td")
    on_buildcmd_file(function (target, batchcmds, sourcefile, opt)
        local llvm_bin = os.getenv("LLVM_DIR") and path.join(os.getenv("LLVM_DIR"), "../../../bin") or "third_party/llvm-install/bin"
        local tblgen = path.join(llvm_bin, is_host("windows") and "mlir-tblgen.exe" or "mlir-tblgen")
        
        if not os.isfile(tblgen) then
            print("Warning: mlir-tblgen not found, skipping TableGen")
            return
        end
        
        local outputdir = path.join(target:autogendir(), "mlir/dialects/tawa")
        local basename = path.basename(sourcefile)
        local include_dirs = {"mlir/dialects", "mlir/dialects/tawa"}
        
        if os.getenv("LLVM_DIR") then
            table.insert(include_dirs, path.join(os.getenv("LLVM_DIR"), "../../../include"))
        end
        
        -- 根据文件名生成不同的输出
        local generators = {}
        if basename == "Dialect" then
            generators = {{"--gen-dialect-decls", "Dialect.h.inc"}, {"--gen-dialect-defs", "Dialect.cpp.inc"}}
        elseif basename == "Ops" then
            generators = {{"--gen-op-decls", "Ops.h.inc"}, {"--gen-op-defs", "Ops.cpp.inc"}}
        elseif basename == "Types" then
            generators = {{"--gen-typedef-decls", "Types.h.inc"}, {"--gen-typedef-defs", "Types.cpp.inc"}}
        elseif basename == "Attributes" then
            generators = {{"--gen-attrdef-decls", "Attributes.h.inc"}, {"--gen-attrdef-defs", "Attributes.cpp.inc"}}
        end
        
        for _, gen in ipairs(generators) do
            local outputfile = path.join(outputdir, gen[2])
            batchcmds:show_progress(opt.progress, "[TableGen] %s", gen[2])
            batchcmds:mkdir(outputdir)
            
            local args = {sourcefile, gen[1], "-o", outputfile}
            for _, inc in ipairs(include_dirs) do
                table.insert(args, "-I")
                table.insert(args, inc)
            end
            
            batchcmds:vrunv(tblgen, args)
            batchcmds:add_depfiles(sourcefile)
        end
    end)

-- Tawa/MLIR 方言库
target("tawa_dialect")
    set_kind("static")
    add_rules("mlir_tablegen")
    add_files("mlir/dialects/tawa/*.td")
    add_files("mlir/dialects/tawa/*.cpp")
    add_includedirs("$(buildir)/_gen", {public = true})
    add_headerfiles("mlir/dialects/tawa/*.td")
    
    -- 链接 MLIR 库（如果可用）
    if os.getenv("LLVM_DIR") or os.isdir("third_party/llvm-install") then
        add_links("MLIRSupport", "LLVMSupport", "LLVMCore")
    end

-- Pass 库
target("tawa_passes")
    set_kind("static")
    add_files("mlir/passes/*.cpp")
    add_deps("tawa_dialect")
    add_includedirs("mlir/dialects", "$(buildir)/_gen")
    
    if os.getenv("LLVM_DIR") or os.isdir("third_party/llvm-install") then
        add_links("MLIRPass", "MLIRTransforms", "MLIRAnalysis")
    end

-- CLI
target("aicompiler")
    set_kind("binary")
    add_files("src/cli/*.cpp")
    add_deps("tawa_dialect", "tawa_passes", "aic_runtime_cpu", "aic_runtime_spirv")
    add_includedirs("mlir/dialects", "$(buildir)/_gen")
    
    if os.getenv("LLVM_DIR") or os.isdir("third_party/llvm-install") then
        add_links("MLIR", "MLIRParser", "MLIRSupport", "LLVMSupport", "LLVMCore")
    end
